# include <opencv2/opencv.hpp>
# include <opencv2/dnn.hpp>
# include "classification.h"
#include "onnxruntime.h"
#include <onnxruntime_cxx_api.h>
using namespace cv;
using namespace std;

using namespace cv::dnn;

Net net;

int softmax(const cv::Mat& src, cv::Mat& dst) {
	float max = 0.0;
	float sum = 0.0;
	max = *std::max_element(src.begin<float>(), src.end<float>());
	cv::exp((src - max), dst);
	sum = cv::sum(dst)[0];
	dst /= sum;
	return 0;
}

//模型加载
bool loadModel(char* modelpath) {
	try {
		net = readNetFromONNX(modelpath);
		return true;
	}
	catch(Exception ex) {
		return false;
	}
}



//单张图像前向传播
void gpuInference(unsigned char* srcPointer, int srcWidth, int srcHeight, int matType, double meanValue, double stdValue, int& labelIndex, double& probability)
{
	try
	{
		Mat srcImage(srcHeight, srcWidth, matType, srcPointer);
		Mat blob, prob;

		//预处理
		resize(srcImage, srcImage, Size(128, 128));
		cvtColor(srcImage, srcImage, cv::COLOR_BGR2RGB);
		srcImage.convertTo(srcImage, CV_32F, 1.0 / 255.0);
		Scalar mean(meanValue, meanValue, meanValue);
		Scalar std(stdValue, stdValue, stdValue);
		subtract(srcImage, mean, srcImage);
		divide(srcImage, std, srcImage);

		//前向传播
		Mat inputBlob = dnn::blobFromImage(srcImage);

		net.setInput(inputBlob);
		prob = net.forward();
		
		//结果处理
		softmax(prob, prob);
		Point maxLoc;
		double maxValue = 0;
		minMaxLoc(prob, 0, &maxValue, 0, &maxLoc);
		labelIndex = maxLoc.x;
		probability = maxValue;
	}
	catch (exception& e)
	{

	}
}






int main() {

	string modelpath = "E:/打包/Inference/Inference/bin/x64/Debug/googlenet.onnx";
	try {
		Mat img = imread("E:/打包/Inference/picture/3.jpg");
		resize(img, img, Size(128, 128));
		cvtColor(img, img, cv::COLOR_BGR2RGB);
		img.convertTo(img, CV_32F, 1.0 / 255.0);
		cv::Scalar default_mean(0.5, 0.5, 0.5);
		cv::Scalar default_std(0.5, 0.5, 0.5);
		cv::subtract(img, default_mean, img);
		cv::divide(img, default_std, img);

		Mat inputBlob = dnn::blobFromImage(img);
		//opencv
		Net net = readNetFromONNX(modelpath);
		clock_t start{ clock() }, end;
		net.setInput(inputBlob);
		Mat prob = net.forward();
		end = clock();
		std::cout << "opencv:"<<end - start<<"ms" << std::endl;
		float* pData = (float*)prob.data;
		for (int i = 0; i < 4; i++) {
			std::cout << *pData++ << std::endl;
		}
		//onnxruntime
		std::wstring model_path(L"E:/打包/Inference/Inference/bin/x64/Debug/googlenet.onnx");
		onnxruntime onr = onnxruntime(model_path, 1, { 1, 3, 128, 128 });
		onr.predict(inputBlob);

		softmax(prob, prob);
		Point maxLoc;
		double maxValue = 0;
		minMaxLoc(prob, 0, &maxValue, 0, &maxLoc);
		int labelIndex = maxLoc.x;
		double probability = maxValue;
		string cla;
		switch (labelIndex)
		{
		case 0:
			cla = "断胶";
			break;
		case 1:
			cla = "多胶";
			break;
		case 2:
			cla = "少胶";
			break;
		case 3:
			cla = "正常";
			break;
		}
		cout << cla << ":" << probability;
	}
	catch (exception ex) {

	}
}